#!/usr/bin/env python2

#       Copyright (C) 2010 SuperSnout

#       This program is free software; you can redistribute it and/or modify
#       it under the terms of the GNU General Public License as published by
#       the Free Software Foundation; either version 2 of the License, or
#       (at your option) any later version.

#       This program is distributed in the hope that it will be useful,
#       but WITHOUT ANY WARRANTY; without even the implied warranty of
#       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#       GNU General Public License for more details.

#       You should have received a copy of the GNU General Public License
#       along with this program; if not, write to the Free Software
#       Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#       MA 02110-1301, USA.

#       Modifies .torrent files


# Use raw_input if we're still using python 2
try:
    input = raw_input
except NameError:
    pass


import argparse
from bencode import bencode, bdecode
import os
import string
import random

parser = argparse.ArgumentParser()

parser.add_argument("-v", "--verbose", action="store_true", dest="verbose", help="switch on verbose mode")
parser.add_argument("-e", "--edit-in-place", action="store_true", dest="replace", help="overwrite files without prompting")
parser.add_argument("-o", "--output", dest="outfile", help="specify output filename. (can not be used if multiple files are given)")
parser.add_argument("-d", "--directory", dest="path", help="specify a directory to save output to")
parser.add_argument("-a", "--announce", dest="url", help="replace announce-url with the one specified")
parser.add_argument("-p", "--private", action="store_true", dest="private", help="make torrent private")
parser.add_argument("-n", "--no-cross-seed", action="store_false", dest="cross", default=True, help="do not randomize info hashes")
parser.add_argument("-c", "--comment", dest="comment", help="replace comment with the one specified")
parser.add_argument("files", nargs='+', help='files to be modified')

args = parser.parse_args()


def random_string(n):
    """Generate a random string of length n"""

    return ''.join(random.choice(string.ascii_letters + string.digits) for i in range(n))


def change_announce(torrent_dict, url):
    """Change announce url."""

    torrent_dict['announce'] = url
    verbose("announce url changed to '%s'" % url)

    if 'announce-list' in torrent_dict:
        del torrent_dict['announce-list']
        verbose("removed all multitracker announce urls")

    return torrent_dict


def change_comment(torrent_dict, comment):
    """Change comment."""

    if not comment:
        if 'comment' in torrent_dict:
            del torrent_dict['comment']
            verbose('removed comment')
    else:
        torrent_dict['comment'] = comment
        verbose("comment changed to '%s'" % comment)

    return torrent_dict


def make_unique(torrent_dict):
    """Add random string to info torrent_dict."""

    torrent_dict['info']['unique'] = uniquestr = random_string(32)
    verbose("random string '%s' added to info torrent_dict" % uniquestr)
    return torrent_dict


def verbose(msg):
    if args.verbose:
        print(msg)


def read_bencode(stream):
    """ Decode a file to an object."""

    with open(stream, "rb") as handle:
        torrent_dict = bdecode(handle.read())

    if torrent_dict:
        verbose("dict read from '%s'" % stream)

    return torrent_dict


def write_bencode(stream, obj):
    """ Encode a given object to a file."""

    with open(stream, "wb") as handle:
        handle.write(bencode(obj))

    verbose("modified dict written to '%s'" % stream)


def confirm(prompt, resp=False):
    """prompts for yes or no response from the user. Returns True for yes and
    False for no.

    'resp' should be set to the default value assumed by the caller when
    user simply types ENTER.

    >>> confirm(prompt='Create Directory?', resp=True)
    Create Directory? Y/n:
    True
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? y/N:
    False
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? y/N: y
    True

    """

    if resp:
        prompt = '%s Y/n: ' % prompt
    else:
        prompt = '%s y/N: ' % prompt

    while True:
        ans = input(prompt)
        if not ans:
            return resp
        if ans not in ['y', 'Y', 'n', 'N']:
            print('please enter y or n.')
            continue
        if ans == 'y' or ans == 'Y':
            return True
        if ans == 'n' or ans == 'N':
            return False


if len(args.files) > 1 and args.outfile:
    parser.error("output filename specified but multiple files given")

for infile in args.files:
    torrent_dict = read_bencode(infile)
    if args.cross:
        torrent_dict = make_unique(torrent_dict)
    if args.url:
        torrent_dict = change_announce(torrent_dict, args.url)
    if args.comment is not None:
        torrent_dict = change_comment(torrent_dict, args.comment)
    if args.private:
        torrent_dict['info']['private'] = 1
        verbose("private set to 1")
    if args.path and args.outfile:
        stream = args.path + args.outfile
        if not os.path.exists(args.path):
            os.makedirs(args.path)
    elif args.path:
        stream = args.path + infile
        if not os.path.exists(args.path):
            os.makedirs(args.path)
    elif args.outfile:
        stream = args.outfile
    else:
        stream = infile
    if args.replace:
        write_bencode(stream, torrent_dict)
    elif os.path.exists(stream) or stream == infile:
        choice = confirm("%s already exists. Overwrite?" % stream, True)
        if choice is not True:
            continue
        else:
            write_bencode(stream, torrent_dict)
    else:
        write_bencode(stream, torrent_dict)
